import math


# 传入明文和列数
def encrypt(msg, col_count):
    msg_len = len(msg)

    # 这个逻辑判读有点多余，可以直接执行下面的逻辑，因为 "*"*0 依然是空字符串。
    # 但是添加这个逻辑判断，或许可以增加代码的可读性
    if msg_len % col_count != 0:
        # 让数据对齐，补齐不够的空白字符
        sub_str = "*"*(col_count-(msg_len % col_count))
        msg += sub_str
        msg_len = len(msg)

    row_count = msg_len//col_count
    result = ""

    for i in range(col_count):
        for j in range(row_count):
            index = i + j*col_count
            result += msg[index]

    return result


if __name__ == '__main__':
    msg = "123456789"
    key = 5
    msg = encrypt(msg, key)
    print("加密后：", msg, "|", len(msg))

    msg = encrypt(msg, math.floor(len(msg)/key))
    print(msg, "|", len(msg))
